{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "class TaskModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(TaskModel, self).__init__()\n",
    "        self.layer1 = nn.Linear(10, 32)  # 假设输入特征是10维的\n",
    "        self.layer2 = nn.Linear(32, 16)\n",
    "        self.output_layer = nn.Linear(16, 1)  # 假设预测目标是1维的\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.layer1(x))\n",
    "        x = F.relu(self.layer2(x))\n",
    "        return self.output_layer(x)\n",
    "\n",
    "# 为两个任务分别实例化模型\n",
    "model1 = TaskModel()\n",
    "model2 = TaskModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.MSELoss()  # 以均方误差为例\n",
    "optimizer1 = optim.Adam(model1.parameters())\n",
    "optimizer2 = optim.Adam(model2.parameters())\n",
    "\n",
    "def soft_parameter_sharing_loss(model1, model2, lambda_):\n",
    "    loss = 0\n",
    "    for param1, param2 in zip(model1.parameters(), model2.parameters()):\n",
    "        loss += torch.sum((param1 - param2) ** 2)\n",
    "    return lambda_ * loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    for inputs, targets1, targets2 in dataloader:  # 假设dataloader加载了数据和标签\n",
    "        optimizer1.zero_grad()\n",
    "        optimizer2.zero_grad()\n",
    "    \n",
    "        outputs1 = model1(inputs)\n",
    "        outputs2 = model2(inputs)\n",
    "    \n",
    "        loss1 = criterion(outputs1, targets1)\n",
    "        loss2 = criterion(outputs2, targets2)\n",
    "    \n",
    "        # 软参数共享正则化项\n",
    "        soft_sharing_loss = soft_parameter_sharing_loss(model1, model2, lambda_=0.01)\n",
    "    \n",
    "        # 总损失\n",
    "        total_loss = loss1 + loss2 + soft_sharing_loss\n",
    "    \n",
    "        total_loss.backward()\n",
    "        optimizer1.step()\n",
    "        optimizer2.step()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
